#! python
# -*- coding: UTF-8 -*-
"""
@Project ：warning 
@File    ：stream_processor.py.py
@IDE     ：PyCharm 
@Author  ：zhangyiheng
@Date    ：2025/2/27 14:35 
@Describe：TODO
"""
import atexit
import cv2
import time
import numpy as np
import torch
from threading import Lock, Thread
from concurrent.futures import ThreadPoolExecutor
from django.core.signals import request_finished
from django.dispatch import receiver
from ultralytics import YOLO
from django.conf import settings

# 初始化模型
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = YOLO("/workspace/lyyj_yolv9e6/weights/best.pt").to(device)
print(f"模型初始化完成 | 运行设备: {device}")

# 全局变量
latest_detections = {}
last_async_sent = {}



class CameraProcessor:
    def __init__(self, url, area_id):
        self.area_id = area_id
        self.url = url
        self.lock = Lock()
        self.frame = None
        self.active = True
        self.last_detection = []
        self.executor = ThreadPoolExecutor(max_workers=8)  # 初始化线程池
        self.status = "disconnected"
        self.last_valid_time = time.time()
        self.error_count = 0

        # 独立资源
        self.cap = None
        self._init_capture()
        Thread(target=self._process_stream, daemon=True).start()

    def _init_capture(self):
        with self.lock:
            try:
                if self.cap and self.cap.isOpened():
                    return

                print(f"初始化摄像头 {self.area_id}...")
                self.cap = cv2.VideoCapture(self.url, cv2.CAP_FFMPEG)
                if self.cap.isOpened():
                    self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
                    self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, settings.CONFIG['frame_size'][0])
                    self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, settings.CONFIG['frame_size'][1])
                    self._update_status("connected")
                else:
                    self._update_status("disconnected")
            except Exception as e:
                print(f"摄像头 {self.area_id} 初始化异常: {str(e)}")

    def _update_status(self, new_status):
        if self.status != new_status:
            self.status = new_status
            print(f"摄像头 {self.area_id} 状态更新: {new_status}")

    def _process_stream(self):
        """处理摄像头流"""
        while self.active:  # 只有在 active 为 True 时才运行
            try:
                # 获取最新帧
                with self.lock:
                    if not self.cap or not self.cap.isOpened():
                        self._init_capture()
                        time.sleep(1)
                        continue

                    # 跳过旧帧
                    for _ in range(2):
                        self.cap.grab()

                    ret, frame = self.cap.retrieve()
                    if ret and frame is not None and frame.size > 0:
                        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                        self.frame = frame
                        self.last_valid_time = time.time()
                        self.error_count = 0
                    else:
                        self.error_count += 1
                        if self.error_count > 3:
                            self._reset_connection()
                        time.sleep(0.1)

                # 异步检测
                if hasattr(self, 'detect_future') and not self.detect_future.done():
                    continue

                if not self.active:  # 确保在退出时不调度新任务
                    break

                self.detect_future = self.executor.submit(
                    self._safe_detect,
                    self.frame.copy() if self.frame is not None else None
                )
                self.detect_future.add_done_callback(
                    lambda f: self._update_detections(f.result())
                )

            except Exception as e:
                print(f"摄像头 {self.area_id} 处理异常: {str(e)}")
                self._reset_connection()
                time.sleep(1)

    def _safe_detect(self, frame):
        if frame is None or not isinstance(frame, np.ndarray) or frame.size == 0:
            return []
        try:
            results = model(frame,
                           imgsz=512,
                           conf=settings.CONFIG['conf_threshold'],
                           device=device,
                           verbose=False,
                           augment=True,
                           max_det=10)
            return [{
                "xyxy": box.xyxy[0].cpu().numpy(),
                "conf": box.conf.item(),
                "cls": int(box.cls.item()),
                "timestamp": time.time()
            } for box in results[0].boxes if box.conf.item() >= settings.CONFIG['conf_threshold']]
        except Exception as e:
            print(f"检测异常 {self.area_id}: {str(e)}")
            return []

    def _update_detections(self, boxes):
        current_time = time.time()
        valid_boxes = [b for b in boxes if current_time - b['timestamp'] < 2]
        self.last_detection = valid_boxes

        # 更新全局状态
        with Lock():
            latest_detections[self.area_id] = [{
                "class": settings.CLASS_CONFIG['names'].get(b["cls"], "未知"),
                "confidence": b["conf"],
                "timestamp": b["timestamp"]
            } for b in valid_boxes]

    def _reset_connection(self):
        with self.lock:
            if self.cap:
                self.cap.release()
            self._init_capture()
            print(f"摄像头 {self.area_id} 已重置")

    def get_current_frame(self):
        with self.lock:
            if self.frame is None or not isinstance(self.frame, np.ndarray):
                return np.zeros((settings.CONFIG['frame_size'][1], settings.CONFIG['frame_size'][0], 3), dtype=np.uint8), False
            return self.frame.copy(), True

    def release(self):
        """释放资源"""
        self.active = False  # 停止主循环
        with self.lock:
            if self.cap and self.cap.isOpened():
                self.cap.release()  # 释放摄像头资源
        self.executor.shutdown(wait=False)  # 关闭线程池


# 初始化摄像头处理器
camera_processors = {aid: CameraProcessor(url, aid) for aid, url in settings.CAMERAS.items()}


def async_result_handler():
    """异步检测结果处理"""
    while True:
        try:
            current_time = time.time()
            with Lock():
                for area_id, dets in latest_detections.items():
                    last_sent = last_async_sent.get(area_id, 0)
                    if current_time - last_sent < settings.CONFIG['async_interval']:
                        continue

                    # 过滤有效检测结果
                    valid_detections = [d for d in dets if (current_time - d["timestamp"]) < 10]
                    if len(valid_detections) >= 2:
                        print(f"[异步处理] 区域 {area_id} 检测到持续目标: {valid_detections}")
                        last_async_sent[area_id] = current_time

            time.sleep(10)  # 每10秒检查一次
        except Exception as e:
            print(f"异步处理异常: {str(e)}")
            time.sleep(5)


def background_cleaner():
    """后台清理线程"""
    while True:
        try:
            time.sleep(settings.CONFIG['cache_clean_interval'])

            # 清理检测数据
            cutoff = time.time() - settings.CONFIG['data_retention']
            with Lock():
                for aid in list(latest_detections.keys()):
                    latest_detections[aid] = [d for d in latest_detections[aid] if d['timestamp'] > cutoff]
                    if not latest_detections[aid]:
                        del latest_detections[aid]

            # 清理GPU缓存
            if device == 'cuda':
                torch.cuda.empty_cache()
                print("已执行GPU缓存清理")
        except Exception as e:
            print(f"后台清理异常: {str(e)}")
            time.sleep(5)


def generate_frames(area_id):
    processor = camera_processors.get(area_id)
    if not processor:
        return

    while processor.active:
        try:
            frame, status_ok = processor.get_current_frame()

            # 绘制检测框
            display_frame = frame.copy()
            current_time = time.time()
            for box in processor.last_detection:
                if current_time - box['timestamp'] < 2:
                    x1, y1, x2, y2 = map(int, box["xyxy"])
                    color = settings.CLASS_CONFIG['colors'].get(box["cls"], (0, 255, 0))
                    cv2.rectangle(display_frame, (x1, y1), (x2, y2), color, 2)
                    label = f"{settings.CLASS_CONFIG['names_en'].get(box['cls'], 'Unknown')} {box['conf']:.2f}"
                    cv2.putText(display_frame, label, (x1, y1 - 10),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)

            # 编码帧
            _, jpeg = cv2.imencode('.jpg', display_frame, [
                cv2.IMWRITE_JPEG_QUALITY, settings.CONFIG['jpeg_quality']
            ])
            yield (b'--frame\r\nContent-Type: image/jpeg\r\n\r\n' + jpeg.tobytes() + b'\r\n')
            time.sleep(1 / settings.CONFIG['max_fps'])

        except Exception as e:
            print(f"视频流 {area_id} 生成异常: {str(e)}")
            time.sleep(0.5)


def cleanup():
    """清理资源"""
    print("\n释放系统资源...")
    for p in camera_processors.values():
        p.release()  # 确保每个摄像头处理器正确释放资源
    if device == 'cuda':
        torch.cuda.empty_cache()


# 注册清理函数
atexit.register(cleanup)

# 使用 Django 信号机制
@receiver(request_finished)
def cleanup_on_request_finished(sender, **kwargs):
    cleanup()


def start_background_threads():
    """启动后台线程"""
    if not hasattr(start_background_threads, '_has_run'):
        Thread(target=async_result_handler, daemon=True).start()
        Thread(target=background_cleaner, daemon=True).start()
        start_background_threads._has_run = True